
#!/usr/bin/env python3
# encoding: utf-8

import os
import sys
import time
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
import yaml

POSITIONS_YAML_PATH = "/home/ubuntu/ros2_ws/src/app/config/positions.yaml"
if os.getenv('ASR_LANGUAGE') == 'Chinese':
    from tweak_tool_ui import Ui_MainWindow
else:
    from tweak_tool_en_ui import Ui_MainWindow
class MainWindow(Ui_MainWindow, QMainWindow):
    def __init__(self):
        super().__init__()
        self.current_func_name = ""

        self.setupUi(self)
        self.pushButton_apply.clicked.connect(self.on_apply_clicked)
        self.pushButton_reset.clicked.connect(self.on_reset_clicked)
        self.pushButton_save.clicked.connect(self.on_save_clicked)
        self.comboBox.currentIndexChanged.connect(self.index_changed)
        self.offset_x.valueChanged.connect(self.on_value_changed)
        self.offset_y.valueChanged.connect(self.on_value_changed)
        self.offset_z.valueChanged.connect(self.on_value_changed)
        self.scale_x.valueChanged.connect(self.on_value_changed)
        self.scale_y.valueChanged.connect(self.on_value_changed)
        self.scale_z.valueChanged.connect(self.on_value_changed)
        for i in range(4):
            getattr(self, "target_" + str(i + 1) + '_x').valueChanged.connect(self.on_value_changed)
            getattr(self, "target_" + str(i + 1) + '_y').valueChanged.connect(self.on_value_changed)
            getattr(self, "target_" + str(i + 1) + '_z').valueChanged.connect(self.on_value_changed)

        self.comboBox.setCurrentIndex(-1)
        self.comboBox.setCurrentIndex(0)
        self.pushButton_apply.setEnabled(False)

    def set_enable_target(self, item_index, enable):
        getattr(self, "target_" + str(item_index) + '_x').setEnabled(enable)
        getattr(self, "target_" + str(item_index) + '_y').setEnabled(enable)
        getattr(self, "target_" + str(item_index) + '_z').setEnabled(enable)
        if not enable:
            getattr(self, "target_" + str(item_index) + '_x').setValue(0)
            getattr(self, "target_" + str(item_index) + '_y').setValue(0)
            getattr(self, "target_" + str(item_index) + '_z').setValue(0)

    def load_params(self):
        with open(POSITIONS_YAML_PATH, 'r') as f:
            params = yaml.safe_load(f)
        return params.get(self.current_func_name, {})

    def get_params_and_set_items(self, target_num):
        params = self.load_params()
        offset = params.get('offset', [0, 0, 0])
        scale = params.get('scale', [1, 1, 1])

        self.offset_x.setValue(offset[0])
        self.offset_y.setValue(offset[1])
        self.offset_z.setValue(offset[2])
        self.scale_x.setValue(scale[0])
        self.scale_y.setValue(scale[1])
        self.scale_z.setValue(scale[2])

        for target_index in range(4):
            self.set_enable_target(target_index + 1, False)

        for target_index in range(target_num):
            self.set_enable_target(target_index + 1, True)
            target = params.get(f'target_{target_index + 1}', [0, 0, 0])
            getattr(self, "target_" + str(target_index + 1) + '_x').setValue(target[0])
            getattr(self, "target_" + str(target_index + 1) + '_y').setValue(target[1])
            getattr(self, "target_" + str(target_index + 1) + '_z').setValue(target[2])

    def save_params(self):
        params = {}
        params['offset'] = [self.offset_x.value(), self.offset_y.value(), self.offset_z.value()]
        params['scale'] = [self.scale_x.value(), self.scale_y.value(), self.scale_z.value()]

        for target_index in range(4):
            widget_x = getattr(self, "target_" + str(target_index + 1) + '_x')
            widget_y = getattr(self, "target_" + str(target_index + 1) + '_y')
            widget_z = getattr(self, "target_" + str(target_index + 1) + '_z')
            if widget_x.isEnabled() and widget_y.isEnabled() and widget_z.isEnabled():
                params[f'target_{target_index + 1}'] = [widget_x.value(), widget_y.value(), widget_z.value()]

        return params

    def index_changed(self, i):
        print("选择了第{}项: {}".format(i, self.comboBox.currentText()))
        if "颜色分拣" == self.comboBox.currentText() or 'Color Sorting' == self.comboBox.currentText():
            self.current_func_name = 'color_sortting'
            self.get_params_and_set_items(3)
        elif "标签分拣" == self.comboBox.currentText() or 'Tag Sorting' == self.comboBox.currentText():
            self.current_func_name = 'tag_sortting'
            self.get_params_and_set_items(3)
        elif "标签码垛" == self.comboBox.currentText() or 'Block Stacking' == self.comboBox.currentText():
            self.current_func_name = 'tag_stackup'
            self.get_params_and_set_items(1)
        elif "垃圾分类" == self.comboBox.currentText() or 'Waste Sorting' == self.comboBox.currentText():
            self.current_func_name = 'waste_classification'
            self.get_params_and_set_items(4)
        else:
            pass
        self.pushButton_apply.setEnabled(False)

    def on_apply_clicked(self):
        params = self.save_params()
        with open(POSITIONS_YAML_PATH, 'r+') as f:
            all_params = yaml.safe_load(f) or {}
            all_params[self.current_func_name] = params
            f.seek(0)
            yaml.dump(all_params, f)
            f.truncate()
        self.pushButton_apply.setEnabled(False)

    def on_reset_clicked(self):
        current_index = self.comboBox.currentIndex()
        self.comboBox.setCurrentIndex(-1)
        self.comboBox.setCurrentIndex(current_index)

    def on_save_clicked(self):
        params = self.save_params()
        with open(POSITIONS_YAML_PATH, 'r+') as f:
            all_params = yaml.safe_load(f) or {}
            all_params[self.current_func_name] = params
            f.seek(0)
            yaml.dump(all_params, f)
            f.truncate()
        self.pushButton_apply.setEnabled(False)

    def on_value_changed(self, value):
        print('change')
        self.pushButton_apply.setEnabled(True)

if __name__ == "__main__":  
    app = QApplication(sys.argv)
    myshow = MainWindow()
    myshow.show()
    sys.exit(app.exec_())

